"""
从emqx归档到tdengine
"""
import json
import logging
import os
import signal
import time
from io import StringIO
import random
from typing import List, Tuple
from datetime import datetime
from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL, CursorResult
from sqlalchemy.sql.elements import TextClause
from concurrent_log_handler import ConcurrentRotatingFileHandler
from paho.mqtt import client as mqtt_client
from paho.mqtt.enums import CallbackAPIVersion

pid = os.getpid()

LOGGER = {
    "LOGGER_NAME": "data_pedestal",
    "TOTAL_LEVEL": logging.INFO,
    "FILE_HANDLER_LEVEL": logging.DEBUG,
    "CONTROL_HANDLER_LEVEL": logging.INFO,
    "FILE_NAME": r"log.txt",
    "MODE": "a",
    "MAX_BYTES": 10 * 1024 * 1024,
    "BACKUP_COUNT": 10,
    "ENCODING": "UTF-8",
    "FORMAT": "%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s"
}


def get_logger() -> logging.getLogger():
    logger = logging.getLogger(LOGGER['LOGGER_NAME'])
    logger.setLevel(LOGGER['TOTAL_LEVEL'])  # 设置日志级别
    file_handler = ConcurrentRotatingFileHandler(filename=LOGGER['FILE_NAME'], mode=LOGGER['MODE'],
                                                 maxBytes=LOGGER['MAX_BYTES'],
                                                 backupCount=LOGGER['BACKUP_COUNT'],
                                                 encoding=LOGGER['ENCODING'])  # 设置输出文件
    file_handler.setLevel(LOGGER['FILE_HANDLER_LEVEL'])
    control_handler = logging.StreamHandler()
    control_handler.setLevel(LOGGER['CONTROL_HANDLER_LEVEL'])
    formatter = logging.Formatter(LOGGER['FORMAT'])
    file_handler.setFormatter(formatter)
    control_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    logger.addHandler(control_handler)
    return logger


# 创建一个实例
logger = get_logger()

try:
    with open('opcua_server_to_td.json', 'r') as f:
        config_txt = f.read()
except FileNotFoundError as e:
    logger.error(str(e))
try:
    config = json.loads(config_txt)
except Exception as e:
    logger.error(str(e))
    os.kill(pid, signal.SIGTERM)

"""
连接 tdengine
"""


class TDEngineWrapper:

    def __init__(self):
        self._engine = create_engine(
            URL.create("taos",
                       username=config["TDENGINE_DATABASE"]["username"],
                       password=config["TDENGINE_DATABASE"]["password"],
                       host=config["TDENGINE_DATABASE"]["host"],
                       port=config["TDENGINE_DATABASE"]["port"],
                       database=config["TDENGINE_DATABASE"]["database"])
        )

    def connected(self):
        try:
            conn = self._engine.connect()
            conn.close()
            return True
        except Exception:
            print(f"和 TDEngine 之间的连接建立失败, 请检查网络是否畅通, 服务器是否正常")
            return False

    def execute(self, statement, parameters=None, *, execution_options=None):
        while not self.connected():
            time.sleep(3)
        with self._engine.begin() as conn:
            if type(statement) is not TextClause:
                statement = text(statement)
            return conn.execute(statement, parameters, execution_options=execution_options)

    @staticmethod
    def rows_to_dict(result: CursorResult, one=False):
        """
        将结果集转成字典
        """
        columns = tuple(map(str, result.keys()))
        rows = result.fetchone() if one else result.fetchall()
        if one:
            return dict(zip(columns, rows))
        return [dict(zip(columns, row)) for row in rows]

    def create_super_table(self,
                           stable: str,
                           ts_field: str,
                           metric_fields: List,
                           tags: List):
        """
        创建超级表
        @stable: 超级表的名称
        @ts_field: 时间戳字段的名称
        @metric_fields: 采集的物理量对应的字段名和类型
        @tags: 标签字段和类型
        """
        if type(metric_fields) is not list:
            raise ValueError("metric_fields 必须是一个列表，里面包含二元组")
        if type(tags) is not list:
            raise ValueError("tags 必须是一个列表，里面包含二元组")
        ddl = f"""
            CREATE STABLE {stable} (
                {ts_field} timestamp,
                {', '.join([' '.join([metric, tp]) for metric, tp in metric_fields])}
            ) TAGS (
                {', '.join([' '.join([tag, tp]) for tag, tp in tags])}
            )
        """
        try:
            self.execute(ddl)
        except Exception as e:
            print(f"超级表建立失败, 错误原因: {e}")
            raise e

    def create_child_table(self, table: str, stable: str, tags_value: Tuple):
        """
        创建子表，必须指定一张超级表
        @table: 子表的名称
        @stable: 超级表的名称
        @tags_value: 标签值, 一个二元组
        """
        ddl = f"""
            CREATE TABLE {table} USING {stable}
            TAGS {tags_value}
        """
        try:
            self.execute(ddl)
        except Exception as e:
            print(f"子表建立失败, 错误原因: {e}")
            raise e

    def td_insert_data(self,
                       table: str,
                       values: List[Tuple],
                       columns: List = None):
        """
        往 tdengine 批量插入数据
        @table: 表名
        @value: 数据
        @columns: 列名
        """
        buf = StringIO()
        for value in values:
            value = list(
                str(item) if isinstance(item, datetime) else "null" if item is None else item
                for item in value
            )
            buf.write(f"({str(value)[1: -1]})")
        # 查看表里面的columns值，获取columns的类型，检验columns的类型

        query = f"""
            INSERT INTO {config["TDENGINE_DATABASE"]["database"]}.{table} {'(' + ','.join(columns) + ')' if columns is not None else ""} 
            VALUES {buf.getvalue()}
        """

        buf.truncate(0)
        buf.close()
        try:
            print(query)
            self.execute(query)
        except Exception as e:
            logger.error(f"数据插入失败, 错误原因: {e}")
            raise e


tdengine_wrapper = TDEngineWrapper()

broker = config['EMQX']['host']
port = config['EMQX']['port']
topic = config['EMQX']['topic_name']
# Generate a Client ID with the subscribe prefix.
client_id = f'subscribe-{random.randint(0, 100)}'


# username = 'emqx'
# password = 'public'


def connect_mqtt() -> mqtt_client:
    def on_connect(client, userdata, flags, rc):
        if rc == 0:
            logger.info("Connected to MQTT Broker!")
            client.subscribe(topic)
        else:
            logger.info("Failed to connect, return code %d\n", rc)

    client = mqtt_client.Client(client_id=client_id, callback_api_version=CallbackAPIVersion.VERSION1)
    # client.username_pw_set(username, password)
    client.on_connect = on_connect
    client.connect(broker, port)
    return client


def subscribe(client: mqtt_client):
    def on_message(client, userdata, msg):
        #print(f"Received `{msg.payload.decode()}` from `{msg.topic}` topic")
        try:
            temp = msg.payload.decode().split(':')
            table_name = temp[0][2:]
            value = temp[-1][:-1]

            str_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            tdengine_wrapper.td_insert_data(table_name, [(str_time, value)], ['ts', 'val'])
        except Exception as e:
            logger.error(e)
    client.subscribe(topic)
    client.on_message = on_message


FIRST_RECONNECT_DELAY = 1
RECONNECT_RATE = 2
MAX_RECONNECT_COUNT = 12
MAX_RECONNECT_DELAY = 60


def on_disconnect(client, userdata, rc):
    logging.info("Disconnected with result code: %s", rc)
    reconnect_count, reconnect_delay = 0, FIRST_RECONNECT_DELAY
    while reconnect_count < MAX_RECONNECT_COUNT:
        logging.info("Reconnecting in %d seconds...", reconnect_delay)
        time.sleep(reconnect_delay)

        try:
            client.reconnect()
            logging.info("Reconnected successfully!")
            return
        except Exception as err:
            logging.error("%s. Reconnect failed. Retrying...", err)

        reconnect_delay *= RECONNECT_RATE
        reconnect_delay = min(reconnect_delay, MAX_RECONNECT_DELAY)
        reconnect_count += 1
    logging.info("Reconnect failed after %s attempts. Exiting...", reconnect_count)


def run():
    client = connect_mqtt()
    subscribe(client)
    client.on_disconnect = on_disconnect
    client.loop_forever()



if __name__ == '__main__':
    run()
